import sys
from pathlib import Path

# Add project root to sys.path to allow running tests from this directory
# The project root is 6 levels up from the parent directory of this file.
project_root = str(Path(__file__).resolve().parents[6])
if project_root not in sys.path:
    sys.path.insert(0, project_root)

#
# Copyright © 2024 Agora
# This file is part of TEN Framework, an open source project.
# Licensed under the Apache License, Version 2.0, with certain conditions.
# Refer to the "LICENSE" file in the root directory for more information.
#
import json
from typing import Any
from unittest.mock import patch, AsyncMock
import tempfile
import os
import asyncio
import filecmp
import shutil
import threading

from ten_runtime import (
    ExtensionTester,
    TenEnvTester,
    Cmd,
    CmdResult,
    StatusCode,
    Data,
)
from ten_ai_base.struct import TTSTextInput, TTSFlush
from ten_ai_base.message import ModuleVendorException, ModuleErrorVendorInfo


# ================ test dump ================
class ExtensionTesterDump(ExtensionTester):
    def __init__(self):
        super().__init__()
        # Use a fixed path as requested by the user.
        self.dump_dir = "./dump/"
        # Use a unique name for the file generated by the test to avoid collision
        # with the file generated by the extension.
        self.test_dump_file_path = os.path.join(
            self.dump_dir, "test_manual_dump.pcm"
        )
        self.audio_end_received = False
        self.received_audio_chunks = []

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        """Called when test starts, sends a TTS request."""
        ten_env_tester.log_info("Dump test started, sending TTS request.")

        tts_input = TTSTextInput(
            request_id="tts_request_1",
            text="hello word, hello agora",
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        if name == "tts_audio_end":
            ten_env.log_info("Received tts_audio_end, stopping test.")
            self.audio_end_received = True
            ten_env.stop_test()

    def on_audio_frame(self, ten_env: TenEnvTester, audio_frame):
        """Receives audio frames and collects their data using the lock/unlock pattern."""
        # The 'audio_frame' object is a wrapper around a memory buffer.
        # We must lock the buffer to safely access the data, copy it,
        # and finally unlock the buffer so the runtime can reuse it.
        buf = audio_frame.lock_buf()
        try:
            # We must copy the data from the buffer, as the underlying memory
            # may be freed or reused after we unlock it.
            copied_data = bytes(buf)
            self.received_audio_chunks.append(copied_data)
        finally:
            # Always ensure the buffer is unlocked, even if an error occurs.
            audio_frame.unlock_buf(buf)

    def write_test_dump_file(self):
        """Writes the collected audio chunks to a file."""
        with open(self.test_dump_file_path, "wb") as f:
            for chunk in self.received_audio_chunks:
                f.write(chunk)

    def find_tts_dump_file(self) -> str | None:
        """Find the dump file created by the TTS extension in the fixed dump directory."""
        if not os.path.exists(self.dump_dir):
            return None
        for filename in os.listdir(self.dump_dir):
            if filename.endswith(".pcm") and filename != os.path.basename(
                self.test_dump_file_path
            ):
                return os.path.join(self.dump_dir, filename)
        return None


@patch("bytedance_tts_duplex.extension.BytedanceV3Client")
def test_dump_functionality(MockBytedanceV3Client):
    """Tests that the dump file from the TTS extension matches the audio received by the test extension."""

    print("Starting test_dump_functionality with mock...")

    # --- Directory Setup ---
    # As requested, use a fixed './dump/' directory.
    DUMP_PATH = "./dump/"

    # Clean up directory before the test, in case of previous failed runs.
    if os.path.exists(DUMP_PATH):
        shutil.rmtree(DUMP_PATH)
    os.makedirs(DUMP_PATH)

    # --- Mock Configuration ---
    mock_instance = MockBytedanceV3Client.return_value
    mock_instance.connect = AsyncMock()
    mock_instance.start_connection = AsyncMock()
    mock_instance.start_session = AsyncMock()
    mock_instance.send_text = AsyncMock()
    mock_instance.finish_session = AsyncMock()
    mock_instance.finish_connection = AsyncMock()
    mock_instance.close = AsyncMock()

    # Create some fake audio data to be streamed
    fake_audio_chunk_1 = b"\x11\x22\x33\x44" * 20
    fake_audio_chunk_2 = b"\xaa\xbb\xcc\xdd" * 20

    # Mock the client constructor to properly handle the response_msgs queue
    def mock_client_init(
        config,
        ten_env,
        vendor,
        response_msgs,
        on_error=None,
        on_usage_characters=None,
        on_fatal_failure=None,
    ):
        # Store the real queue passed by the extension
        mock_instance.response_msgs = response_msgs

        # Populate the queue with mock data asynchronously
        async def populate_queue():
            # Use constants directly
            EVENT_TTSResponse = 352
            EVENT_SessionFinished = 152

            await asyncio.sleep(0.01)  # Small delay to let the extension start
            await response_msgs.put((EVENT_TTSResponse, fake_audio_chunk_1))
            await asyncio.sleep(0.01)
            await response_msgs.put((EVENT_TTSResponse, fake_audio_chunk_2))
            await asyncio.sleep(0.01)
            await response_msgs.put((EVENT_SessionFinished, b""))

        # Start the population task
        asyncio.create_task(populate_queue())
        return mock_instance

    MockBytedanceV3Client.side_effect = mock_client_init

    # --- Test Setup ---
    tester = ExtensionTesterDump()

    dump_config = {
        "params": {
            "app_id": "valid_appid_for_test",
            "token": "valid_token_for_test",
        },
        "appid": "valid_appid_for_test",
        "token": "valid_token_for_test",
        "dump": True,
        "dump_path": DUMP_PATH,
    }

    tester.set_test_mode_single("bytedance_tts_duplex", json.dumps(dump_config))

    try:
        print("Running dump test...")
        tester.run()
        print("Dump test completed.")

        # --- Assertions ---
        assert tester.audio_end_received, "tts_audio_end was not received"

        # Write the audio chunks collected by the test extension to its own dump file
        tester.write_test_dump_file()
        assert os.path.exists(
            tester.test_dump_file_path
        ), "Test dump file was not created"

        # Find the dump file automatically created by the TTS extension
        tts_dump_file = tester.find_tts_dump_file()
        assert (
            tts_dump_file is not None
        ), f"Could not find TTS-generated dump file in {DUMP_PATH}"

        print(f"Comparing TTS dump file: {tts_dump_file}")
        print(f"With test dump file:    {tester.test_dump_file_path}")

        # Binary comparison of the two files
        assert filecmp.cmp(
            tts_dump_file, tester.test_dump_file_path, shallow=False
        ), "The TTS dump file and the test-generated dump file do not match."

        print("✅ Dump file binary comparison passed.")

    finally:
        # Cleanup the dump directory after the test.
        if os.path.exists(DUMP_PATH):
            shutil.rmtree(DUMP_PATH)


# ================ test flush ================
class ExtensionTesterFlush(ExtensionTester):
    def __init__(self):
        super().__init__()
        self.ten_env: TenEnvTester | None = None
        self.audio_start_received = False
        self.first_audio_frame_received = False
        self.flush_start_received = False
        self.audio_end_received = False
        self.flush_end_received = False
        self.audio_end_reason = ""
        self.total_audio_duration_from_event = 0
        self.received_audio_bytes = 0
        self.sample_rate = 24000
        self.bytes_per_sample = 2  # 16-bit
        self.channels = 1
        self.audio_received_after_flush_end = False

    def on_start(self, ten_env_tester: TenEnvTester) -> None:
        self.ten_env = ten_env_tester
        ten_env_tester.log_info("Flush test started, sending long TTS request.")
        tts_input = TTSTextInput(
            request_id="tts_request_for_flush",
            text="This is a very long text designed to generate a continuous stream of audio, providing enough time to send a flush command.",
        )
        data = Data.create("tts_text_input")
        data.set_property_from_json(None, tts_input.model_dump_json())
        ten_env_tester.send_data(data)
        ten_env_tester.on_start_done()

    def on_audio_frame(self, ten_env: TenEnvTester, audio_frame):
        if self.flush_end_received:
            ten_env.log_error("Received audio frame after tts_flush_end!")
            self.audio_received_after_flush_end = True

        if not self.first_audio_frame_received:
            self.first_audio_frame_received = True
            ten_env.log_info("First audio frame received, sending flush data.")
            flush_data = Data.create("tts_flush")
            flush_data.set_property_from_json(
                None,
                TTSFlush(flush_id="tts_request_for_flush").model_dump_json(),
            )
            ten_env.send_data(flush_data)

        buf = audio_frame.lock_buf()
        try:
            self.received_audio_bytes += len(buf)
        finally:
            audio_frame.unlock_buf(buf)

    def on_data(self, ten_env: TenEnvTester, data) -> None:
        name = data.get_name()
        ten_env.log_info(f"on_data name: {name}")

        if name == "tts_audio_start":
            self.audio_start_received = True
            return

        if name == "tts_flush_start":
            self.flush_start_received = True
            return

        json_str, _ = data.get_property_to_json(None)
        if not json_str:
            return
        payload = json.loads(json_str)
        ten_env.log_info(f"on_data payload: {payload}")

        if name == "tts_audio_end":
            self.audio_end_received = True
            self.audio_end_reason = payload.get("reason")
            self.total_audio_duration_from_event = payload.get(
                "request_total_audio_duration_ms"
            )

        elif name == "tts_flush_end":
            self.flush_end_received = True

            def stop_test_later():
                ten_env.log_info("Waited after flush_end, stopping test now.")
                ten_env.stop_test()

            # Use threading.Timer to allow a short grace period to catch stray audio frames
            timer = threading.Timer(0.5, stop_test_later)
            timer.start()

    def get_calculated_audio_duration_ms(self) -> int:
        duration_sec = self.received_audio_bytes / (
            self.sample_rate * self.bytes_per_sample * self.channels
        )
        return int(duration_sec * 1000)


@patch("bytedance_tts_duplex.extension.BytedanceV3Client")
def test_flush_logic(MockBytedanceV3Client):
    """
    Tests that sending a flush command during TTS streaming correctly stops
    the audio and sends the appropriate events.
    """
    print("Starting test_flush_logic with mock...")

    # --- Mock Configuration ---
    mock_instance = MockBytedanceV3Client.return_value
    mock_instance.connect = AsyncMock()
    mock_instance.start_connection = AsyncMock()
    mock_instance.start_session = AsyncMock()
    mock_instance.send_text = AsyncMock()
    mock_instance.finish_session = AsyncMock()
    mock_instance.finish_connection = AsyncMock()
    mock_instance.close = AsyncMock()

    # Create a cancel event to signal the mock audio stream to stop
    cancel_event = asyncio.Event()

    # When flush is called in the extension, it should trigger this cancel method
    async def mock_cancel():
        cancel_event.set()

    mock_instance.cancel = AsyncMock(side_effect=mock_cancel)

    # Mock the client constructor
    def mock_client_init(
        config,
        ten_env,
        vendor,
        response_msgs,
        on_error=None,
        on_usage_characters=None,
        on_fatal_failure=None,
    ):
        mock_instance.response_msgs = response_msgs

        async def populate_queue():
            EVENT_TTSMetric = 888
            EVENT_TTSResponse = 352
            EVENT_SessionFinished = 152

            await asyncio.sleep(0.5)

            # Continuously send audio chunks until cancelled
            for _ in range(20):
                if cancel_event.is_set():
                    # bytedance doesn't have a specific flush event from the client,
                    # the flush is handled by stopping the session.
                    await response_msgs.put((EVENT_SessionFinished, b""))
                    return

                # send metric event
                await response_msgs.put((EVENT_TTSMetric, b""))

                # send audio chunk
                await response_msgs.put(
                    (EVENT_TTSResponse, b"\x11\x22\x33" * 100)
                )
                await asyncio.sleep(0.1)

            # This part is only reached if not cancelled
            await response_msgs.put((EVENT_SessionFinished, b""))

        asyncio.create_task(populate_queue())
        return mock_instance

    MockBytedanceV3Client.side_effect = mock_client_init

    # --- Test Setup ---
    config = {
        "appid": "a_valid_appid",
        "token": "a_valid_token",
        "params": {
            "app_id": "a_valid_appid",
            "token": "a_valid_token",
        },
    }
    tester = ExtensionTesterFlush()
    tester.set_test_mode_single("bytedance_tts_duplex", json.dumps(config))

    print("Running flush logic test...")
    tester.run()
    print("Flush logic test completed.")

    # --- Assertions ---
    assert tester.audio_start_received, "Did not receive tts_audio_start."
    assert tester.first_audio_frame_received, "Did not receive any audio frame."
    assert tester.audio_end_received, "Did not receive tts_audio_end."
    assert tester.flush_end_received, "Did not receive tts_flush_end."

    # In bytedance, a flushed stream ends with 'flush' reason
    assert (
        tester.audio_end_reason == 2
    ), f"Expected audio end reason 'flush', but got '{tester.audio_end_reason}'"

    calculated_duration = tester.get_calculated_audio_duration_ms()
    event_duration = tester.total_audio_duration_from_event
    print(
        f"Calculated duration: {calculated_duration}ms, Event duration: {event_duration}ms"
    )

    print("✅ Flush logic test passed successfully.")
